#include "SDCF.h"
#include "Static.h"
#include "CFG.h"
#include "call_graph.h"

/* 
 * global variables anouncement 
 */

RTSTATUS sdcf_status(false, false, true, true, false);     //sdcf runtime status

file_t f;                            //log file
file_t cg_f;                         //cg log file
uint client_id;						 //client id
CallGraph* global_cg;                //global call graph

/*
 * events need registering 
 */

static void event_exit(void);
static dr_emit_flags_t event_basic_block(void *drcontext, void *tag, instrlist_t *bb, bool for_trace, bool translating);

/* 
 * init procedure 
 */

DR_EXPORT void 
dr_init(client_id_t id)
{
	/* variables initializing */
	client_id = id;
	global_cg = new CallGraph;

	/* log file initializing */
	char logname[512];
	char *dirsep;
	int len;

	len = dr_snprintf(logname, sizeof(logname)/sizeof(logname[0]),"%s", dr_get_client_path(id));
	DR_ASSERT(len > 0);
	for (dirsep = logname + len; *dirsep != DIRSEP; dirsep--)
		DR_ASSERT(dirsep > logname);
	
	len = dr_snprintf(dirsep + 1,
		(sizeof(logname) - (dirsep - logname))/sizeof(logname[0]),
		"sys.dot");
	DR_ASSERT(len > 0);
	logname[sizeof(logname)/sizeof(logname[0])-1] = '\0';
	f = dr_open_file(logname, DR_FILE_WRITE_OVERWRITE);
	DR_ASSERT(f != INVALID_FILE);

	len = dr_snprintf(dirsep + 1,
		(sizeof(logname) - (dirsep - logname))/sizeof(logname[0]),
		"cg.dot");
	DR_ASSERT(len > 0);
	logname[sizeof(logname)/sizeof(logname[0])-1] = '\0';
	cg_f = dr_open_file(logname, DR_FILE_WRITE_OVERWRITE);
	DR_ASSERT(cg_f != INVALID_FILE);

	dr_fprintf(cg_f,"digraph CG {\n");

	/* register events */

	dr_register_bb_event(event_basic_block);
	dr_register_exit_event(event_exit);

}// end of dr_init

/* 
 *end event 
 */

static void 
event_exit(void)
{	
	dr_fprintf(cg_f,"}\n");
	dr_fprintf(f,"ends\n");

	dr_close_file(f);
	dr_close_file(cg_f);

	delete global_cg;
}// end of event_exit

/* 
 * bb entered event 
 */

static dr_emit_flags_t
event_basic_block(void *drcontext, void *tag, instrlist_t *bb,
				  bool for_trace, bool translating)
{

	/* local variables anouncement */
	instr_t* first_instruction;
	instr_t* last_instruction;
	int last_opcode;
	app_pc first_pc;
	app_pc last_pc;
	module_data_t* module;

	static app_pc expected_call;                //the expected call address for current call
	static app_pc temp_expected_ret;            //temporuary expected return
	static int temp_cfg_id;

	/* get current information */
	
	first_instruction = instrlist_first(bb);           //first instruction
	first_pc = instr_get_app_pc(first_instruction);    //first pc
	module = dr_lookup_module(first_pc);               //current module
	last_instruction = instrlist_last(bb);             //last instrction
	last_pc = instr_get_app_pc(last_instruction);
	last_opcode = instr_get_opcode(last_instruction);  //last opcode

	/* module procedure */
    if(module){
		Get_Mod_Info(module);
	}

	if(!sdcf_status.in_mod_exe){
		if(sdcf_status.init){
			sdcf_status.init = false;
		}
		return DR_EMIT_DEFAULT;
	}
	
	BB_Print(drcontext, bb);
   
	/* if it's the first bb in exe */
	if(sdcf_status.init_exe){                                        
		sdcf_status.init_exe = false;		
		function_node *temp_func_node = global_cg->createFunctionNode((uint)first_pc, 0,CFG_FILE_PREFIX);

		global_cg->setHeadFunction(temp_func_node);
		global_cg->current_func_node = temp_func_node;
		dr_fprintf(f,"find root for cg\n");
	}

   /* cg construction procedure*/
	if(sdcf_status.just_called){
		sdcf_status.just_called = false;

		if(expected_call - first_pc == 0){                     //which means the call actually takes place

			dr_fprintf(f,"expected_call met: 0x%08x\n", expected_call);
			function_node *temp_func_node = global_cg->createFunctionNode((uint)expected_call, (uint)temp_expected_ret, "");		
			global_cg->addCallee(global_cg->current_func_node, temp_func_node);

			dr_snprintf(temp_func_node->function_name, 30, "%s%d", CFG_FILE_PREFIX, temp_cfg_id++);
			
			dr_fprintf(cg_f,"\"%s 0x%08x\" -> \"%s 0x%08x\"\n",global_cg->current_func_node->function_name,
				global_cg->current_func_node->entrance_address,
				temp_func_node->function_name,temp_func_node->entrance_address);
			dr_fprintf(f,"\"0x%08x\" -> \"0x%08x\"\n",global_cg->current_func_node->entrance_address, temp_func_node->entrance_address);

			global_cg->current_func_node = temp_func_node;
			dr_fprintf(f,"call added\n");
		}
	}

	if(sdcf_status.just_returned){
		sdcf_status.just_returned = false;
		
		if(global_cg->current_func_node->expected_return_address - (uint)first_pc == 0){                      //which means the ret actually takes place
			dr_fprintf(f,"expected_ret met, 0x%08x\n", global_cg->current_func_node->expected_return_address);

			StaticCompletion(global_cg->current_func_node->cfg);
			global_cg->printSingleNode(global_cg->current_func_node);
			global_cg->current_func_node = global_cg->current_func_node->caller;
		}
	}

	/* add bb */
	if(!global_cg->current_func_node->cfg->endCFG)
	{
		basicblock_t *newBB = create_basicblock();
		newBB->first_address = (uint)first_pc;
		newBB->last_address = (uint)last_pc;
		newBB->last_length = (uint)instr_length(drcontext,last_instruction);
		newBB->next_branch = NULL;
		newBB->isStatic = false;
		instruction_list_t *tmpList = create_instruction_list();
		for(instr_t* currentInstr = instrlist_first(bb);currentInstr!=NULL;currentInstr=instr_get_next(currentInstr)){
			instruction_t *tmpInstr = create_instruction(false,(uint)instr_get_app_pc(currentInstr));
			tmpInstr->instr = currentInstr;
			instruction_list_append(tmpList,tmpInstr);
		}
		newBB->bb_instr_list = tmpList;

		///////// add the new BB to CFG /////////// 
		newBB = global_cg->current_func_node->cfg->addBB(newBB);
		if(newBB!=NULL)
		{
			//dr_fprintf(f," @add BB %d at 0x%08x to CFG\n",newBB->id,newBB->first_address);
			if(global_cg->current_func_node->cfg->getSize()==1)
			{
				global_cg->current_func_node->cfg->setEntrance(newBB);
			}
			if(global_cg->current_func_node->cfg->linkedBB!=NULL)
			{
				global_cg->current_func_node->cfg->linkBB(global_cg->current_func_node->cfg->linkedBB,newBB,false);
				//dr_fprintf(f," @link BB %d to %d\n",linkedBB->id,newBB->id);
			}
			global_cg->current_func_node->cfg->linkedBB = newBB;
		}
		else
		{
			dr_fprintf(f," @Error when add BB to CFG\n");
		}

		/////////// add branch point to the list ///////////
		if(global_cg->current_func_node->cfg->meetBranch){
			bool branchpoint_branched = false;
			if(global_cg->current_func_node->cfg->followingAddress == (uint)instr_get_app_pc(first_instruction))
				branchpoint_branched = false;
			else branchpoint_branched = true;
			global_cg->current_func_node->cfg->addBranchPoint(global_cg->current_func_node->cfg->branchAddress,
				global_cg->current_func_node->cfg->followingAddress,
				global_cg->current_func_node->cfg->branchedAddress,branchpoint_branched);
			global_cg->current_func_node->cfg->meetBranch = false;
			dr_fprintf(f," @add branch point to CFG: branchAddress: 0x%08x\n", global_cg->current_func_node->cfg->branchAddress);
		}
		if((last_opcode >= 26 && last_opcode <=41) || (last_opcode >= 152 && last_opcode <= 167))
		{
			//////// indirect branch cannot be handled now ///////
			if(instr_reads_memory(last_instruction))
			{
			}
			else{
				global_cg->current_func_node->cfg->meetBranch = true;
				global_cg->current_func_node->cfg->branchAddress = (uint)instr_get_app_pc(last_instruction);
				global_cg->current_func_node->cfg->branchedAddress = (uint)opnd_get_pc(instr_get_target(last_instruction));
				global_cg->current_func_node->cfg->followingAddress = global_cg->current_func_node->cfg->branchAddress
					+(uint)instr_length(drcontext,last_instruction);
			}
		}
	}

	/* cti procedure */
	if(instr_is_call(last_instruction)){
		sdcf_status.just_called = true;
		if(opnd_is_pc(instr_get_target(last_instruction)))
		{
			expected_call = opnd_get_pc(instr_get_target(last_instruction));
			dr_fprintf(f,"expected_call: 0x%08x\n", expected_call);
			temp_expected_ret = last_pc + instr_length(drcontext, last_instruction);
			dr_fprintf(f,"temp_expected_ret: 0x%08x\n", temp_expected_ret);
			//dr_fprintf(f,"call: 0x%08x\n", last_pc + instr_length(drcontext, last_instruction));
		}
	}
	else if(instr_is_return(last_instruction)){	
		sdcf_status.just_returned = true;
	}

	if(sdcf_status.init){
		sdcf_status.init = false;
	}

	return DR_EMIT_DEFAULT;
}// end of event_basic_block